
#include "stdafx.h"

#include <string>

#include <vcclr.h>

#include "ImageCodecPlugin.h"

using namespace std;
using namespace Imf;
using namespace Derm;

using namespace System::Runtime::InteropServices;

/*===========================================================================*/

/*===========================================================================*/

/// OpenEXR input stream
class ExrInputStream : public IStream
{
public:

	ExrInputStream(Stream ^io) :
		IStream(""),
		_io (io)
	{
		
	}

	virtual ~ExrInputStream()
	{
		
	}

	virtual bool read(char c[], int n)
	{
		cli::array<unsigned char> ^buffer = gcnew cli::array<unsigned char>(n + 2);
		int red;

		// Read from stream
		red = _io->Read(buffer, 0, n);
		// Marshall read data into 'c'
		Marshal::Copy(buffer, 0, (IntPtr)c, n);

		return (red == n);
	}

	virtual Imf::Int64 tellg()
	{
		return (_io->Position);
	}

	virtual void seekg(Imf::Int64 pos)
	{
		_io->Seek(pos, SeekOrigin::Begin);
	}

	virtual void clear() { };

private:

	gcroot<Stream ^> _io;

};

class ExrOutputStream : public OStream
{
public:

	ExrOutputStream(Stream ^io) :
		OStream(""),
		_io (io)
	{
			
	}

	virtual ~ExrOutputStream()
	{
		
	}

	virtual void write (const char c[], int n)
	{
		cli::array< Byte >^ buffer = gcnew array< Byte >(n + 2);

		_io->Write(buffer, 0, n);
	}

	virtual Imf::Int64 tellp()
	{
		return (_io->Position);
	}

	virtual void seekp(Imf::Int64 pos)
	{
		_io->Seek(pos, SeekOrigin::Current);
	}

private:

	gcroot<Stream ^> _io;

};


/*===========================================================================*/

Pixel::Type OpenEXRImageCodecPlugin::GetPixelTypeFromChannels(const Imf::ChannelList &channels)
{
	const Imf::Channel *rChannel = channels.findChannel("R");
	const Imf::Channel *gChannel = channels.findChannel("G");
	const Imf::Channel *bChannel = channels.findChannel("B");
	const Imf::Channel *aChannel = channels.findChannel("A");
	const Imf::Channel *yChannel = channels.findChannel("Y");
	const Imf::Channel *byChannel = channels.findChannel("BY");
	const Imf::Channel *ryChannel = channels.findChannel("RY");

	if        (rChannel && gChannel && bChannel && aChannel) {
		// Check whether all components have the same pixel type
		if ((rChannel->type != gChannel->type) || (rChannel->type != bChannel->type) || (rChannel->type != aChannel->type))
			throw gcnew InvalidOperationException("RGBA channels doesn't have the same data type");
		switch(rChannel->type) {
			case Imf::UINT:
				throw gcnew InvalidOperationException("RGBA UINT channels are not supported");
				break;
			case Imf::HALF:
				return (Pixel::Type::RGBAHF);
			case Imf::FLOAT:
				return (Pixel::Type::RGBAF);
			default:
				throw gcnew InvalidOperationException("RGBA channels data type is unknown");
		}
	} else if (rChannel && gChannel && bChannel) {
		// Check whether all components have the same pixel type
		if ((rChannel->type != gChannel->type) || (rChannel->type != bChannel->type))
			throw gcnew InvalidOperationException("RGB channels doesn't have the same data type");
		switch(rChannel->type) {
			case Imf::UINT:
				throw gcnew InvalidOperationException("RGB UINT channels are not supported");
				break;
			case Imf::HALF:
				return (Pixel::Type::RGBHF);
			case Imf::FLOAT:
				return (Pixel::Type::RGBF);
			default:
				throw gcnew InvalidOperationException("RGB channels data type is unknown");
		}
	} else if (yChannel && ryChannel && byChannel && aChannel) {
		switch(yChannel->type) {
			case Imf::UINT:
				throw gcnew InvalidOperationException("YCbCrA UINT channels are not supported");
				break;
			case Imf::HALF:
				return (Pixel::Type::BGRAHF);
			case Imf::FLOAT:
				return (Pixel::Type::BGRAF);
			default:
				throw gcnew InvalidOperationException("YCbCrA channels data type is unknown");
		}
	} else if (yChannel && ryChannel && byChannel) {
		switch(yChannel->type) {
			case Imf::UINT:
				throw gcnew InvalidOperationException("YCbCr UINT channels are not supported");
				break;
			case Imf::HALF:
				return (Pixel::Type::BGRAHF);
			case Imf::FLOAT:
				return (Pixel::Type::BGRAF);
			default:
				throw gcnew InvalidOperationException("YCbCr channels data type is unknown");
		}
	} else if (yChannel) {
		switch(yChannel->type) {
			case Imf::UINT:
				throw gcnew InvalidOperationException("GRAY UINT channels are not supported");
				break;
			case Imf::HALF:
				return (Pixel::Type::GRAYHF);
			case Imf::FLOAT:
				return (Pixel::Type::GRAYF);
			default:
				throw gcnew InvalidOperationException("GRAY channels data type is unknown");
		}
	} else if (gChannel) {
		switch(gChannel->type) {
			case Imf::UINT:
				throw gcnew InvalidOperationException("GREEN UINT channels are not supported");
				break;
			case Imf::HALF:
				return (Pixel::Type::GRAYHF);
			case Imf::FLOAT:
				return (Pixel::Type::GRAYF);
			default:
				throw gcnew InvalidOperationException("GREEN channels data type is unknown");
		}
	} else {
		string layerNames;

		for (Imf::ChannelList::ConstIterator it = channels.begin(), ite = channels.end(); it != ite; it++)
			layerNames.append(it.name());
			
		throw gcnew InvalidOperationException(String::Format("unknown channels configuration {0}", gcnew String(layerNames.c_str())));
	}
}

bool OpenEXRImageCodecPlugin::CheckAvailability()
{
	return (true);
}

ImageInfo ^OpenEXRImageCodecPlugin::QueryImageInfo(Stream ^fs, ImageCodec::CodecCriteria ^criteria)
{
	ImageInfo ^info = gcnew ImageInfo();
	ExrInputStream stream = ExrInputStream(fs);

	Imf::InputFile file(stream, 1);
	
	// Plugin support only EXR format
	info->ContainerFormat = ImageFormat::Exr;

	// Get image size information
	const Imath::Box2i &dataWindow = file.header().dataWindow();
	info->Width = dataWindow.max.x - dataWindow.min.x + 1;
	info->Height = dataWindow.max.y - dataWindow.min.y + 1;

	// Get image pixel type information
	info->PixelType = GetPixelTypeFromChannels(file.header().channels());

	return (info);
}

Image ^OpenEXRImageCodecPlugin::Load(Stream ^fs, ImageCodec::CodecCriteria ^criteria)
{
	Image ^image;
	ExrInputStream stream = ExrInputStream(fs);
	Imf::InputFile file(stream, 1);
	const Imf::ChannelList &channels = file.header().channels();
	Pixel::Type pType;

	// Create image
	const Imath::Box2i &dataWindow = file.header().dataWindow();
	pType = GetPixelTypeFromChannels(file.header().channels());
	image = Image::CreateImage(dataWindow.max.x - dataWindow.min.x + 1, dataWindow.max.y - dataWindow.min.y + 1, pType);

	// Extract pixel data
	Imf::FrameBuffer frameBuffer;
	size_t offset = - dataWindow.min.x * image->PixelSize - dataWindow.min.y * image->Stride;

	const Imf::Channel *rChannel = channels.findChannel("R");
	const Imf::Channel *gChannel = channels.findChannel("G");
	const Imf::Channel *bChannel = channels.findChannel("B");
	const Imf::Channel *aChannel = channels.findChannel("A");
	const Imf::Channel *yChannel = channels.findChannel("Y");
	const Imf::Channel *byChannel = channels.findChannel("RY");
	const Imf::Channel *ryChannel = channels.findChannel("BY");

	System::Object ^imageData = image->GetPixelData();
	GCHandle ^hImageData = GCHandle::Alloc(imageData, GCHandleType::Pinned);
	char *imageDataPtr = reinterpret_cast<char*>(hImageData->AddrOfPinnedObject().ToPointer());

	if        (rChannel && gChannel && bChannel) {
		size_t imagePixelComponentSize = image->PixelSize / Pixel::GetComponentsCount(image->PixelFormat);

		frameBuffer.insert(
			"R",
			Imf::Slice(
				rChannel->type, (char*)(imageDataPtr + 0 * imagePixelComponentSize + offset),
				/* (X Stride) */ image->PixelSize,
				/* (Y Stride) */ image->Stride, /* X/Y Sampling */ 1, 1
			)
		);
		frameBuffer.insert(
			"G",
			Imf::Slice(
				gChannel->type, (char*)(imageDataPtr + 1 * imagePixelComponentSize + offset),
				/* (X Stride) */ image->PixelSize,
				/* (Y Stride) */ image->Stride, /* X/Y Sampling */ 1, 1
			)
		);
		frameBuffer.insert(
			"B",
			Imf::Slice(
				bChannel->type, (char*)(imageDataPtr + 2 * imagePixelComponentSize + offset),
				/* (X Stride) */ image->PixelSize,
				/* (Y Stride) */ image->Stride, /* X/Y Sampling */ 1, 1
			)
		);
		if (aChannel) {
			frameBuffer.insert(
				"A",
				Imf::Slice(
					aChannel->type, (char*)(imageDataPtr + 3 * imagePixelComponentSize + offset),
					/* (X Stride) */ image->PixelSize,
					/* (Y Stride) */ image->Stride, /* X/Y Sampling */ 1, 1
				)
			);
		}

		// Load pixels
		file.setFrameBuffer(frameBuffer);
		file.readPixels(dataWindow.min.y, dataWindow.max.y);

	} else if (yChannel && ryChannel && byChannel) {
		stream.seekg(0);
		Imf::RgbaInputFile rgbaFile(stream);

		// Read YCbCr(A) EXR using RGBA interface

		// Read the entire RGB(A)
		rgbaFile.setFrameBuffer((Imf::Rgba*)imageDataPtr - dataWindow.min.x - dataWindow.min.y * image->Width, 1, image->Width);
		rgbaFile.readPixels(dataWindow.min.y, dataWindow.max.y);
	} else if (yChannel) {
		frameBuffer.insert(
			"Y",
			Imf::Slice(
				yChannel->type, (char*)(imageDataPtr + offset),
				/* (X Stride) */ image->PixelSize,
				/* (Y Stride) */ image->Stride, /* X/Y Sampling */ 1, 1
			)
		);

		// Load pixels
		file.setFrameBuffer(frameBuffer);
		file.readPixels(dataWindow.min.y, dataWindow.max.y);
	} else if (gChannel) {
		frameBuffer.insert(
			"G",
			Imf::Slice(
				gChannel->type, (char*)(imageDataPtr + offset),
				/* (X Stride) */ image->PixelSize,
				/* (Y Stride) */ image->Stride, /* X/Y Sampling */ 1, 1
			)
		);

		// Load pixels
		file.setFrameBuffer(frameBuffer);
		file.readPixels(dataWindow.min.y, dataWindow.max.y);
	}

	hImageData->Free();

	// OpenEXR specify image lines from top to bottom
	image->FlipVertically();

	return (image);
}

void OpenEXRImageCodecPlugin::Save(Stream ^fs, Image ^image, ImageFormat format, ImageCodec::CodecCriteria ^criteria)
{
	try {
		// Check allowed pixel formats
		switch (image->PixelFormat) {
			case Pixel::Type::GRAYF:
			case Pixel::Type::RGBF:
			case Pixel::Type::RGBHF:
			case Pixel::Type::RGBAF:
			case Pixel::Type::RGBAHF:
				/* Allowed pixel format */
				break;
			default:
				throw gcnew ArgumentException("image pixel type " + image->PixelFormat.ToString() + " not supported", "image");
		}

		Imf::Compression exrCompress = Imf::PIZ_COMPRESSION;		// Default compression method

		// Determine compression parameter
		if (criteria->IsDefined("Compression") == true) {
			String ^compressionMethod = (String ^)criteria["Compression"];

			if        (String::Compare(compressionMethod, "None", true) == 0) {
				exrCompress = Imf::NO_COMPRESSION;
			} else if (String::Compare(compressionMethod, "Zip", true) == 0) {
				exrCompress = Imf::ZIP_COMPRESSION;
			} else if (String::Compare(compressionMethod, "Piz", true) == 0) {
				exrCompress = Imf::PIZ_COMPRESSION;
			} else if (String::Compare(compressionMethod, "Pxr24", true) == 0) {
				exrCompress = Imf::PXR24_COMPRESSION;
			} else if (String::Compare(compressionMethod, "ExrB44", true) == 0) {
				exrCompress = Imf::B44_COMPRESSION;
			} else
				throw gcnew ArgumentException("invalid value '" + compressionMethod + "' for codec criteria 'Compression'");
		}

		int width  = image->Width;
		int height = image->Height;
		int dx = 0, dy = 0;

		Imath::Box2i dataWindow(Imath::V2i(0, 0), Imath::V2i(width - 1, height - 1));
		Imath::Box2i displayWindow (Imath::V2i(-dx, -dy), Imath::V2i(width - dx - 1, height - dy - 1));

		Imf::Header header = Imf::Header(displayWindow, dataWindow, 1, Imath::V2f(0,0), 1, Imf::INCREASING_Y, exrCompress);
		Imf::PixelType exrPixelType = Imf::HALF;

		// Keep precision of single-precision floating-point components
		switch (image->PixelFormat) {
			case Pixel::Type::RGBF:
			case Pixel::Type::RGBAF:
				exrPixelType = Imf::FLOAT;
				break;
		}

		switch (image->PixelFormat) {
			case Pixel::Type::GRAYF:
				header.channels().insert ("Y", Imf::Channel(exrPixelType));
				break;
			case Pixel::Type::RGBF:
			case Pixel::Type::RGBHF:
				header.channels().insert ("R", Imf::Channel(exrPixelType));
				header.channels().insert ("G", Imf::Channel(exrPixelType));
				header.channels().insert ("B", Imf::Channel(exrPixelType));
				break;
			case Pixel::Type::RGBAF:
			case Pixel::Type::RGBAHF:
				header.channels().insert ("R", Imf::Channel(exrPixelType));
				header.channels().insert ("G", Imf::Channel(exrPixelType));
				header.channels().insert ("B", Imf::Channel(exrPixelType));
				header.channels().insert ("A", Imf::Channel(exrPixelType));
				break;
		}

		// Setup pixel data framebuffer
		Imf::FrameBuffer frameBuffer;

		GCHandle ^hImageData = GCHandle::Alloc(image->GetPixelData(), GCHandleType::Pinned);
		char *imageDataPtr = reinterpret_cast<char*>(hImageData->AddrOfPinnedObject().ToPointer());
		size_t imagePixelComponentSize = image->PixelSize / Pixel::GetComponentsCount(image->PixelFormat);

		switch (image->PixelFormat) {
			case Pixel::Type::GRAYF:
				frameBuffer.insert ("Y", Imf::Slice (exrPixelType, imageDataPtr, image->PixelSize, image->Stride));
				break;
			case Pixel::Type::RGBF:
			case Pixel::Type::RGBHF:
				frameBuffer.insert ("R", Imf::Slice (exrPixelType, imageDataPtr + 0 * imagePixelComponentSize, image->PixelSize, image->Stride));
				frameBuffer.insert ("G", Imf::Slice (exrPixelType, imageDataPtr + 1 * imagePixelComponentSize, image->PixelSize, image->Stride));
				frameBuffer.insert ("B", Imf::Slice (exrPixelType, imageDataPtr + 2 * imagePixelComponentSize, image->PixelSize, image->Stride));
				break;
			case Pixel::Type::RGBAF:
			case Pixel::Type::RGBAHF:
				frameBuffer.insert ("R", Imf::Slice (exrPixelType, imageDataPtr + 0 * imagePixelComponentSize, image->PixelSize, image->Stride));
				frameBuffer.insert ("G", Imf::Slice (exrPixelType, imageDataPtr + 1 * imagePixelComponentSize, image->PixelSize, image->Stride));
				frameBuffer.insert ("B", Imf::Slice (exrPixelType, imageDataPtr + 2 * imagePixelComponentSize, image->PixelSize, image->Stride));
				frameBuffer.insert ("A", Imf::Slice (exrPixelType, imageDataPtr + 3 * imagePixelComponentSize, image->PixelSize, image->Stride));
				break;
		}

		// Write pixel data
		ExrOutputStream stream = ExrOutputStream(fs);
		Imf::OutputFile file(stream, header);

		file.setFrameBuffer(frameBuffer);
		file.writePixels(height);

		hImageData->Free();
	} catch(Iex::BaseExc &e) {

	}
}
